from typing import Dict, List, Optional, Any
from uuid import uuid4
from datetime import datetime
from loguru import logger
from pydantic import BaseModel

class Task(BaseModel):
    task_id: str
    description: str
    status: str
    created_at: datetime
    completed_at: Optional[datetime] = None
    result: Optional[Any] = None

class Agent:
    def __init__(self):
        self.tasks: Dict[str, Task] = {}
        self.tools: Dict[str, callable] = {}
        logger.info("Agent initialized")

    def register_tool(self, name: str, tool_func: callable):
        """Register a new tool that the agent can use."""
        self.tools[name] = tool_func
        logger.info(f"Registered new tool: {name}")

    def get_available_tools(self) -> List[str]:
        """Get list of available tools."""
        return list(self.tools.keys())

    def create_task(self, description: str) -> Task:
        """Create a new task."""
        task_id = str(uuid4())
        task = Task(
            task_id=task_id,
            description=description,
            status="pending",
            created_at=datetime.now()
        )
        self.tasks[task_id] = task
        logger.info(f"Created new task: {task_id}")
        return task

    def get_task(self, task_id: str) -> Optional[Task]:
        """Get task by ID."""
        return self.tasks.get(task_id)

    def process_task(self, task_id: str) -> Optional[Task]:
        """Process a task and return the result."""
        task = self.get_task(task_id)
        if not task:
            logger.error(f"Task not found: {task_id}")
            return None

        try:
            task.status = "processing"
            # Here you would implement the actual task processing logic
            # This could involve using the registered tools, making API calls,
            # or performing any other necessary operations
            
            # For demonstration, we'll just echo the task description
            task.result = f"Processed task: {task.description}"
            task.status = "completed"
            task.completed_at = datetime.now()
            logger.info(f"Task completed: {task_id}")
        except Exception as e:
            task.status = "failed"
            task.result = str(e)
            logger.error(f"Task failed: {task_id}, error: {e}")

        return task 